# TODO: 导入必要的库和模块
from PIL import Image, ImageDraw
import numpy as np
import gradio as gr
import pickle

# TODO: 加载保存的KNN模型，这样我们可以使用预训练的模型进行预测
model_path = 'best_knn_model.pkl'
with open(model_path, 'rb') as file:
    knn_model = pickle.load(file)
# TODO: 定义预测函数，这个函数将用于Gradio接口进行预测
def predict(stroke_data):
    try:
        # 将 Sketchpad 数据转换为图像
        image = Image.new("RGB", (280, 280), "white")
        draw = ImageDraw.Draw(image)
        for stroke in stroke_data:
            # 确保坐标点是整数类型的元组
            points = [(int(float(point[0]) + (280 / 2)), int(float(point[1]) + (280 / 2))) for point in stroke]
            draw.line(points, fill="black", width=4)
        image = image.resize((28, 28), Image.Resampling.LANCZOS)

        # 转换为灰度
        image = image.convert('L')

        # 将图像数据归一化到 [0, 1] 范围内
        image = np.array(image) / 255.0

        # 展平图像数据为一维数组
        features = image.flatten()

        # 使用模型进行预测
        prediction = knn_model.predict([features])
        return prediction[0]
    except Exception as e:
        print(f"Error processing stroke data: {e}")
        print(f"Stroke data received: {stroke_data}")
        return "Error processing your drawing."
# TODO: 创建Gradio接口，这个接口将用于用户输入和显示预测结果
def create_interface(model):
    with gr.Blocks() as demo:
        gr.Markdown("# Handwritten Digit Recognition")
        with gr.Row():
            with gr.Column(scale=3):
                gr.Markdown("## Draw Digit")
                sketchpad = gr.Sketchpad(label="Handwritten Digit", width=280, height=280)
            with gr.Column(scale=1):
                gr.Markdown("## Prediction")
                prediction = gr.Textbox(label="Prediction")
        with gr.Row():
            submit = gr.Button("Predict")
            submit.click(predict, inputs=sketchpad, outputs=prediction)

    return demo

# TODO: 启动Gradio接口，用户可以通过这个接口进行交互
interface = create_interface(knn_model)
interface.launch(share=True)
